from django.shortcuts import render
import hanlp
import nltk

# 加载预训练模型进行
# 分词任务
tok = hanlp.load(hanlp.pretrained.tok.COARSE_ELECTRA_SMALL_ZH)
# 命名实体识别任务
ner = hanlp.load(hanlp.pretrained.ner.MSRA_NER_ELECTRA_SMALL_ZH)
# 确保你下载了句子分割模型
nltk.download('punkt')
nltk.download('punkt_tab')


def entity_insight_view(request):
    initial_example = "2024年，中国的GDP增长率预计将达到5.5%。\n华为在2023年获得了100亿元的融资，进一步扩展其5G网络业务。" # 初始例句
    user_input = initial_example  # 用于保存用户输入的文本
    entities = []
    highlighted_text = ""  # 初始化高亮文本
    show_result = False

    # 标签到 CSS 类的映射
    tag_to_class = {
        'PERSON': 'entity-person',
        'LOCATION': 'entity-location',
        'ORGANIZATION': 'entity-organization',
        # 金融有关标签
        'COMPANY_NAME': 'entity-company',
        'MONEY': 'entity-money',
        'DATE': 'entity-date',
        'TRANSACTION_ID': 'entity-transaction',
        'PERCENT': 'entity-percent'
    }

    if request.method == 'POST':
        show_result = True
        # 获取用户输入的文本
        text = request.POST['text']
        user_input = text

        # 按行分割文本，以保留空行
        lines = text.split('\n')

        for line in lines:
            if line.strip() == "":
                highlighted_text += "<br>"  # 添加空行
                continue

            # 使用 nltk 进行句子分割
            sentences = nltk.sent_tokenize(line)

            for sentence in sentences:
                text_tok = tok(sentence)  # 分词
                ner_results = ner([text_tok])  # 获取命名实体识别结果

                # 处理 NER 结果
                for entity in ner_results[0]:
                    word, tag, start, end = entity
                    css_class = tag_to_class.get(tag, 'entity-unknown')
                    # 在实体后面添加标签
                    sentence = sentence.replace(word,
                                                f'<span class="{css_class}" data-ner="{tag}">{word} ({tag})</span>')

                # 将处理后的句子添加到高亮文本中
                highlighted_text += sentence + "<br>"  # 添加换行

                # 收集所有实体
                for entity in ner_results[0]:
                    word, tag, start, end = entity
                    entities.append((word, tag))

    return render(request, 'entity_insight/index.html', {
        'highlighted_text': highlighted_text,
        'entities': entities,
        'user_input': user_input,  # 将用户输入传回模板
        'show_result': show_result,
    })
